import pylab as plt
import os
from params import *
from fit import *

dataset = 'gowalla'
file = os.path.join(WORKDIR, 'gowalla', 'temporal', "results","%s_edge_gaps.txt"%dataset)

#DEGREES = [2,3,4,5,6,7,8,9,10,20,30,40,50,60,70,80,90,100]
#DEGREES = [2]
#DEGREES = range(358,1001)
DEGREES = range(2,100)

a_values = []
b_values = []
mle = {}
for T in DEGREES:
    hist = {}
    values = []

    for line in open(file):
        degree, gap = map(int,line.strip().split())
        if degree == T and gap > 0:
            values.append(gap)

    if len(values) < 2:
      continue
    avg = float(sum(values)) / len(values)

    stor = ModelSelector(values)
    models = stor.get_models()
    baseline = 0
    for model in models:
        ll = model['Log-likelihood']
        name = model['Name']
        print name, model['Label']
        if name == 'Exponential':
            baseline = ll
            xfit4, yfit4 = model['Fit']
        if name == 'Shifted exponential':
          xfit,yfit = model['Fit']
        if name == 'Truncated power-law':
          c, a, b = model['Params']
          xfit2,yfit2 = model['Fit']
        if name == 'Power-law':
          xfit3,yfit3 = model['Fit']

    mle[T] = {}
    for model in models:
      mle[T][model['Name']] = 100.0*(baseline - model['Log-likelihood'])/baseline

    if T == DEGREES[0]:
        plt.figure()
        plt.clf()
        plt.axes(FIG_AXES2)
        x,y = stor.data()
        plt.loglog(xfit2,yfit2,'k--',linewidth=1)
        plt.loglog(xfit,yfit,'k:',linewidth=1)
        plt.loglog(xfit3,yfit3,'k-',linewidth=1)
        #plt.loglog(xfit4,yfit4,'k.',linewidth=1)
        #plt.legend(['Trunc. power-law', 'Shifted exp.', 'Power-law', 'Exp.'],
        #    loc='lower left')
        plt.legend(['Trunc. power-law', 'Shifted exp.', 'Power-law'],
            loc='lower left')
        plt.loglog(x,y,'ko',mfc="None")
        plt.xlim((0,100))
        plt.xlabel(r"$\delta(%d)$ [days]"%(T-1))
        plt.ylabel('PMF')
        plt.grid(True)
        plt.savefig("gap_prob_%d.pdf"%(T-1))
        plt.close()
        break

